import os
import time
from threading import Thread, Event
from models import TaskStatus, TaskType

# 任务类，单例模式
class Task:
    _instance = None
    _subprocess = None

    def __new__(cls, *args, **kwargs):
        if cls._instance is not None:
            cls._terminate_subprocess()
        cls._instance = super(Task, cls).__new__(cls)
        return cls._instance

    def __init__(self, taskID, model, task, strategyType, method, lr, epoch, batchSize, dataStartTime, dataEndTime):
        if not hasattr(self, 'initialized'):  # 防止重复初始化
            self.taskID = taskID
            self.model = model
            self.taskType = task
            self.strategyType = strategyType
            self.method = method
            self.lr = lr
            self.epoch = epoch
            self.batchSize = batchSize
            self.dataStartTime = dataStartTime
            self.dataEndTime = dataEndTime
            self.initialized = True
            self.log_file = None
            self.losses = []
            self.accuracies = []
            self.status = TaskStatus.NOT_STARTED
            self.stop_event = Event()  # 添加停止事件
            self.thread = None  # 添加线程引用

    def start(self, is_fake=True):
        # 取文件路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        static_dir = os.path.join(current_dir, "static")
        
        # 确保 static 目录存在
        os.makedirs(static_dir, exist_ok=True)
        
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        self.log_file = os.path.join(static_dir, f"loss_{timestamp}.log")

        if is_fake:
            # 启动假任务
            def fake_task_thread():
                if self.strategyType == TaskType.COLLABORATIVE_EVOLUTION:
                    image_loss_file = os.path.join(static_dir, "image-loss-co.txt")
                else:
                    image_loss_file = os.path.join(static_dir, "image-loss.txt")
                print(f"Image loss file path: {image_loss_file}")
                
                if not os.path.exists(image_loss_file):
                    print(f"Error: {image_loss_file} does not exist!")
                    return
                
                try:
                    with open(image_loss_file, "r") as src_file:
                        with open(self.log_file, "w", buffering=1) as dst_file:  # 使用行缓冲
                            for line in src_file:
                                if self.stop_event.is_set():  # 检查是否需要停止
                                    self.status = TaskStatus.TERMINATED
                                    return
                                dst_file.write(line)
                                dst_file.flush()  # 强制写入
                                os.fsync(dst_file.fileno())  # 确保写入磁盘
                                time.sleep(5)  # 模拟训练过程的时间间隔
                    self.status = TaskStatus.COMPLETED
                except Exception as e:
                    print(f"Error during file operations: {str(e)}")
                    self.status = TaskStatus.ERROR

            self.status = TaskStatus.RUNNING
            self.stop_event.clear()  # 重置停止事件
            self.thread = Thread(target=fake_task_thread)
            self.thread.daemon = True
            self.thread.start()
        else:
            self.status = TaskStatus.RUNNING
            # TODO: 启动真实训练任务
            
            self.status = TaskStatus.COMPLETED
            pass

    def get_metrics(self):
        if not os.path.exists(self.log_file):
            print(f"Log file does not exist: {self.log_file}")
            return [0] * 10, [0] * 10
        
        try:
            self.losses = []
            self.accuracies = []
            with open(self.log_file, 'r') as f:
                for line in f:
                    # 检查是否包含 "Loss" 和 "Accuracy"
                    if "Loss:" in line and "Accuracy:" in line:
                        try:
                            # 提取 Loss 值
                            loss_start = line.find("Loss:") + 5
                            loss_end = line.find(" ", loss_start)
                            loss = float(line[loss_start:loss_end])
                            
                            # 提取 Accuracy 值
                            acc_start = line.find("Accuracy:") + 9
                            acc_end = line.find("\n", acc_start)
                            if acc_end == -1:  # 如果是最后一行
                                acc_end = len(line)
                            acc = float(line[acc_start:acc_end])
                            
                            self.losses.append(loss)
                            self.accuracies.append(acc)
                        except (ValueError, IndexError) as e:
                            print(f"Error parsing line: {line.strip()}, Error: {str(e)}")
                            continue
        except Exception as e:
            print(f"Error reading metrics: {str(e)}")
            return [0] * 10, [0] * 10
        
        # 确保列表长度为10，不足的用0补充
        self.losses = (self.losses + [0] * 10)[:10]
        self.accuracies = (self.accuracies + [0] * 10)[:10]
        
        return self.losses, self.accuracies
    
    @classmethod
    def _terminate_subprocess(cls):
        if cls._subprocess is not None:
            cls._subprocess.terminate()
            cls._subprocess = None

    def terminate(self):
        """终止当前任务"""
        if self.status == TaskStatus.RUNNING:
            self.stop_event.set()  # 设置停止事件
            if self.thread and self.thread.is_alive():
                self.thread.join(timeout=2)  # 等待线程结束，最多等待2秒
            self._terminate_subprocess()  # 终止子进程
            self.status = TaskStatus.TERMINATED

if __name__ == "__main__":
    task = Task("model", "task", "method", 0.001, 10, 32, "2024-01-01", "2024-01-02")
    task.start(is_fake=True)
    
    # 每秒读取一次数据
    for _ in range(15):  # 测试15秒
        losses, accuracies = task.get_metrics()
        print(f"Current metrics - Losses: {losses}, Accuracies: {accuracies}")
        time.sleep(1)
